#!/usr/bin/python
# -*- coding: UTF-8 -*-
"""
Environment is a 2D car.
Car has 5 sensors to obtain distance information.
Car collision => reward = -1, otherwise => reward = 0.
You can train this RL by using LOAD = False, after training, this model will be store in the a local folder.
Using LOAD = True to reload the trained model for playing.
You can customize this script in a way you want.
View more on [莫烦Python] : https://morvanzhou.github.io/tutorials/
Requirement:
pyglet >= 1.2.4
numpy >= 1.12.1
tensorflow >= 1.0.1
"""
import tensorflow as tf
import numpy as np
import os
import shutil  # 高层次的文件操作工具
from car_env import CarEnv
'''
seed( ) 用于指定随机数生成时所用算法开始的整数值。
1.如果使用相同的seed( )值，则每次生成的随即数都相同；
2.如果不设置这个值，则系统根据时间来自己选择这个值，此时每次生成的随机数因时间差异而不同。
3.设置的seed()值仅一次有效
'''
np.random.seed(1)
tf.set_random_seed(1)

MAX_EPISODES = 500
MAX_EP_STEPS = 600

LR_A = 1e-4  # learning rate for actor
LR_C = 1e-4  # learning rate for critic
GAMMA = 0.9  # reward discount

REPLACE_ITER_A = 800
REPLACE_ITER_C = 700

MEMORY_CAPACITY = 2000
BATCH_SIZE = 16
VAR_MIN = 0.1
RENDER = True
LOAD = False
DISCRETE_ACTION = False
# 建立environment
env = CarEnv(discrete_action=DISCRETE_ACTION)
STATE_DIM = env.state_dim  # 5
ACTION_DIM = env.action_dim  # 1
ACTION_BOUND = env.action_bound  # [-1,1]
# all placeholder for tf  # name_scope不会作为tf.get_variable变量的前缀，但是会作为tf.Variable的前缀
with tf.name_scope('S'):
    S = tf.placeholder(tf.float32, shape=[None, STATE_DIM], name='s')
with tf.name_scope('R'):
    R = tf.placeholder(tf.float32, [None, 1], name='r')
with tf.name_scope('S_'):
    S_ = tf.placeholder(tf.float32, shape=[None, STATE_DIM], name='s_')


# 演员
class Actor(object):
    # 有了__init__方法，在创建实例的时候，就不能传入空的参数了，必须传入与__init__方法匹配的参数，
    # 但self不需要传，Python解释器自己会把实例变量传进去
    def __init__(self, sess, action_dim, action_bound, learning_rate,
                 t_replace_iter):
        # self 相当于c++的this指针
        self.sess = sess
        self.a_dim = action_dim
        self.action_bound = action_bound
        self.lr = learning_rate
        self.t_replace_iter = t_replace_iter
        self.t_replace_counter = 0

        with tf.variable_scope('Actor'):
            # input s, output a
            self.a = self._build_net(S, scope='eval_net',
                                     trainable=True)  # actor当前网络，需要训练，用于选取动作
            # input s_, output a, get a_ for critic
            self.a_ = self._build_net(
                S_, scope='target_net',
                trainable=False)  # actor目标网络，无需训练，用于计算Q_target

        # tf.get_collection获取训练变量等效用法
        # GLOBAL_VARIABLES: 该collection默认加入所有的Variable对象，并且在分布式环境中共享
        # 参考：tf.global_variables.通常,所有TRAINABLE_VARIABLES 变量都将在 MODEL_VARIABLES,
        # 所有 MODEL_VARIABLES 变量都将在 GLOBAL_VARIABLES
        self.e_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                          scope='Actor/eval_net')
        self.t_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
                                          scope='Actor/target_net')

    def _build_net(self, s, scope, trainable):
        with tf.variable_scope(scope):
            # 初始化权重的初始化程序'Xavier',这个初始化器是用来保持每一层的梯度大小都差不多相同
            init_w = tf.contrib.layers.xavier_initializer()
            # 生成一个初始值为常量value的tensor对象
            init_b = tf.constant_initializer(0.001)
            # 全链接层
            net = tf.layers.dense(s,
                                  100,
                                  activation=tf.nn.relu,
                                  kernel_initializer=init_w,
                                  bias_initializer=init_b,
                                  name='l1',
                                  trainable=trainable)
            net = tf.layers.dense(net,
                                  20,
                                  activation=tf.nn.relu,
                                  kernel_initializer=init_w,
                                  bias_initializer=init_b,
                                  name='l2',
                                  trainable=trainable)
            with tf.variable_scope('a'):
                actions = tf.layers.dense(net,
                                          self.a_dim,
                                          activation=tf.nn.tanh,
                                          kernel_initializer=init_w,
                                          name='a',
                                          trainable=trainable)
                # tf.multiply（）两个矩阵中对应元素各自相乘   # self.action_bound = [-1,1]
                scaled_a = tf.multiply(
                    actions, self.action_bound, name='scaled_a'
                )  # Scale output to -action_bound to action_bound
        return scaled_a

    def learn(self, s):  # batch update
        self.sess.run(self.train_op, feed_dict={S: s})
        # 如果T%C=1,则更新Critic目标网络和Actor目标网络参数
        if self.t_replace_counter % self.t_replace_iter == 0:
            # tf.assign() 赋值
            self.sess.run([
                tf.assign(t, e) for t, e in zip(self.t_params, self.e_params)
            ])
        self.t_replace_counter += 1

    def choose_action(self, s):
        s = s[np.newaxis, :]  # single state
        return self.sess.run(self.a, feed_dict={S: s})[0]  # actor当前网络，用于选取动作

    def add_grad_to_graph(self, a_grads):
        with tf.variable_scope('policy_grads'):
            # grad_ys 是ys的加权向量列表，和ys长度相同  # 这是从 Critic 来的 dQ/da
            # tf.gradients()实现ys对xs求导
            # actor的梯度分为两部分
            self.policy_grads = tf.gradients(ys=self.a,
                                             xs=self.e_params,
                                             grad_ys=a_grads)

        with tf.variable_scope('A_train'):
            opt = tf.train.RMSPropOptimizer(
                -self.lr)  # (- learning rate) for actor
            # zip() 函数用于将可迭代的对象作为参数，将对象中对应的元素打包成一个个元组，然后返回由这些元组组成的列表
            # minimize()实际上包含了两个步骤，即compute_gradients和apply_gradients，前者用于计算梯度，后者用于使用计算得到的梯度来更新对应的variable
            # 通过神经网络的梯度反向传播来更新Actor当前网络的所有参数𝜃
            # apply_gradients(grads_and_vars, global_step=None, name=None)
            # grads_and_vars: List of (gradient, variable) #list包括的是：梯度和更新变量的元组对  # 执行对应变量的更新梯度操作
            self.train_op = opt.apply_gradients(
                zip(self.policy_grads, self.e_params))  # self.e_params:actor


# 评论家
class Critic(object):
    def __init__(self, sess, state_dim, action_dim, learning_rate, gamma,
                 t_replace_iter, a, a_):
        self.sess = sess
        self.s_dim = state_dim
        self.a_dim = action_dim
        self.lr = learning_rate
        self.gamma = gamma
        self.t_replace_iter = t_replace_iter
        self.t_replace_counter = 0
        self.time_step = 0

        with tf.variable_scope('Critic'):
            # Input (s, a), output q
            self.a = a
            self.q = self._build_net(S, self.a, 'eval_net', trainable=True)

            # Input (s_, a_), output q_ for q_target
            # 𝜋𝜃′(𝜙(𝑆′𝑗))是通过Actor目标网络得到 a_
            self.q_ = self._build_net(
                S_, a_, 'target_net', trainable=False
            )  # target_q is based on a_ from Actor's target_net # actor目标网络
            self.e_params = tf.get_collection(
                tf.GraphKeys.GLOBAL_VARIABLES,
                scope='Critic/eval_net')  # critic当前网络，需要训练
            self.t_params = tf.get_collection(
                tf.GraphKeys.GLOBAL_VARIABLES,
                scope='Critic/target_net')  # crtic_目标网络，不需要训练

        with tf.variable_scope('target_q'):
            # 𝑄′(𝜙(𝑆′𝑗),𝜋𝜃′(𝜙(𝑆′𝑗)),𝑤′)则是通过Critic目标网络得到的
            self.target_q = R + self.gamma * self.q_
        # TD_error
        with tf.variable_scope('TD_error'):
            self.loss = tf.reduce_mean(
                tf.squared_difference(self.target_q, self.q))
            tf.summary.scalar('loss', self.loss)

        with tf.variable_scope('C_train'):
            # 使用均方差损失函数，通过神经网络的梯度反向传播来更新Critic当前网络的所有参数𝑤
            self.train_op = tf.train.RMSPropOptimizer(self.lr).minimize(
                self.loss)

        with tf.variable_scope('a_grad'):
            self.a_grads = tf.gradients(
                self.q,
                a)[0]  # tensor of gradients of each sample (None, a_dim)

    def _build_net(self, s, a, scope, trainable):
        with tf.variable_scope(scope):
            init_w = tf.contrib.layers.xavier_initializer()
            init_b = tf.constant_initializer(0.01)

            with tf.variable_scope('l1'):
                n_l1 = 100
                w1_s = tf.get_variable('w1_s', [self.s_dim, n_l1],
                                       initializer=init_w,
                                       trainable=trainable)
                w1_a = tf.get_variable('w1_a', [self.a_dim, n_l1],
                                       initializer=init_w,
                                       trainable=trainable)
                b1 = tf.get_variable('b1', [1, n_l1],
                                     initializer=init_b,
                                     trainable=trainable)
                net = tf.nn.relu6(tf.matmul(s, w1_s) + tf.matmul(a, w1_a) + b1)
            net = tf.layers.dense(net,
                                  20,
                                  activation=tf.nn.relu,
                                  kernel_initializer=init_w,
                                  bias_initializer=init_b,
                                  name='l2',
                                  trainable=trainable)
            with tf.variable_scope('q'):
                q = tf.layers.dense(net,
                                    1,
                                    kernel_initializer=init_w,
                                    bias_initializer=init_b,
                                    trainable=trainable)  # Q(s,a)
        return q

    def learn(self, s, a, r, s_, t_):
        self.time_step += 1
        self.sess.run(self.train_op, feed_dict={S: s, self.a: a, R: r, S_: s_})

        rs = sess.run(merged, feed_dict={S: s, self.a: a, R: r, S_: s_})
        writer.add_summary(rs, self.time_step)

        # 如果T%C=1,则更新Critic目标网络和Actor目标网络参数
        if self.t_replace_counter % self.t_replace_iter == 0:
            self.sess.run([
                tf.assign(t, e) for t, e in zip(self.t_params, self.e_params)
            ])
        self.t_replace_counter += 1


class Memory(object):
    def __init__(self, capacity, dims):
        self.capacity = capacity
        self.data = np.zeros((capacity, dims))
        self.pointer = 0

    def store_transition(self, s, a, r, s_):
        transition = np.hstack((s, a, [r], s_))
        index = self.pointer % self.capacity  # replace the old memory with new memory
        self.data[index, :] = transition
        self.pointer += 1

    def sample(self, n):
        assert self.pointer >= self.capacity, 'Memory has not been fulfilled'
        indices = np.random.choice(self.capacity, size=n)
        return self.data[indices, :]


sess = tf.Session()
# Create actor and critic.
actor = Actor(sess, ACTION_DIM, ACTION_BOUND[1], LR_A, REPLACE_ITER_A)
critic = Critic(sess, STATE_DIM, ACTION_DIM, LR_C, GAMMA, REPLACE_ITER_C,
                actor.a, actor.a_)
actor.add_grad_to_graph(critic.a_grads)  # dQ/da

M = Memory(MEMORY_CAPACITY, dims=2 * STATE_DIM + ACTION_DIM + 1)

saver = tf.train.Saver()
path = './discrete' if DISCRETE_ACTION else './continuous'

if LOAD:
    saver.restore(sess, tf.train.latest_checkpoint(path))  # ???????????
else:
    sess.run(tf.global_variables_initializer())

merged = tf.summary.merge_all()
writer = tf.summary.FileWriter("./logs", sess.graph)


def train():
    var = 2.  # control exploration
    for ep in range(MAX_EPISODES):
        s = env.reset()
        ep_step = 0
        for t in range(MAX_EP_STEPS):
            # while True:
            if RENDER:
                env.render()
            # Added exploration noise   # actor当前网络选取动作
            a = actor.choose_action(s)
            a = np.clip(np.random.normal(a, var), *ACTION_BOUND
                        )  # add randomness to action selection for exploration
            # 执行动作𝐴,得到新状态𝑆′,奖励𝑅,是否终止状态
            s_, r, done = env.step(a)
            M.store_transition(s, a, r, s_)
            if M.pointer > MEMORY_CAPACITY:
                var = max([var * .9995,
                           VAR_MIN])  # decay the action randomness
                # 从经验回放集合𝐷中采样𝑚个样本
                b_M = M.sample(BATCH_SIZE)
                b_s = b_M[:, :STATE_DIM]
                b_a = b_M[:, STATE_DIM:STATE_DIM + ACTION_DIM]
                b_r = b_M[:, -STATE_DIM - 1:-STATE_DIM]
                b_s_ = b_M[:, -STATE_DIM:]

                critic.learn(b_s, b_a, b_r, b_s_, t)
                actor.learn(b_s)

            s = s_
            ep_step += 1

            if done or t == MAX_EP_STEPS - 1:
                # if done:
                print(
                    'Ep:',
                    ep,
                    '| Steps: %i' % int(ep_step),
                    '| Explore: %.2f' % var,
                )
                break

    # 判断路径是否为目录  # 递归地删除文件夹以及里面的文件
    if os.path.isdir(path):
        shutil.rmtree(path)
    os.mkdir(path)
    ckpt_path = os.path.join(path, 'DDPG.ckpt')
    save_path = saver.save(sess, ckpt_path, write_meta_graph=False)
    print("\nSave Model %s\n" % save_path)


def eval():
    env.set_fps(30)
    while True:
        s = env.reset()
        while True:
            env.render()
            a = actor.choose_action(s)
            s_, r, done = env.step(a)
            s = s_
            if done:
                break


if __name__ == '__main__':
    if LOAD:
        eval()
    else:
        train()
